from flask import Flask, request
import json
import downloadCivitai
import downloadLiblib

app = Flask(__name__)


@app.route('/')
def hello():
    return 'Hello, World!'


@app.route('/fetchModelWeb', methods=['POST'])
def fetchModelWeb():
    """提取图片的主体"""
    requestJson = request.json
    webUrl = requestJson.get('webUrl', "")  # 是否裁切
    id = requestJson.get('id', None)  # 是否裁切
    tabIndex = requestJson.get('tabIndex', 0)  # 是否裁切

    try:
        if "civitai.com" in webUrl:
            imgUrl, localImgPath, nameOrigin, type, version, triggerWord, downloadUrl = downloadCivitai.startFetch(
                webUrl)
            res = {"webUrl": webUrl, "id": id, "imgUrl": imgUrl, "localImgPath": localImgPath,
                   "nameOrigin": nameOrigin, "type": type,
                   "version": version, "triggerWord": triggerWord, "downloadUrl": downloadUrl}
            return json.dumps(res, indent=2)
        elif "www.liblib" in webUrl or "liblib.ai" in webUrl or "liblib.art" in webUrl or "liblib.com" in webUrl:
            imgUrl, localImgPath, nameOrigin, type, version, triggerWord, downloadUrl = downloadLiblib.startFetch(
                webUrl, tabIndex=tabIndex)
            res = {"webUrl": webUrl, "id": id, "imgUrl": imgUrl, "localImgPath": localImgPath,
                   "nameOrigin": nameOrigin, "type": type,
                   "version": version, "triggerWord": triggerWord, "downloadUrl": downloadUrl}
            return json.dumps(res, indent=2)
    except Exception as e:
        print(e)
        res = {"webUrl": webUrl}
        return json.dumps(res, indent=2)

    res = {"webUrl": webUrl}
    return json.dumps(res, indent=2)


if __name__ == '__main__':
    app.run(host="0.0.0.0", port=20102)
